package middleware

import (
	"fmt"
	"net/http"
	"strconv"
	"sync"
	"time"

	"gitee.com/zhenyangze/gin-framework/internal/app/bases"
	"gitee.com/zhenyangze/gin-framework/internal/app/providers"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis"
)

// RateLimiter provides advanced rate limiting capabilities
type RateLimiter struct {
	redis        *redis.Client
	windowSize   time.Duration
	maxRequests  int64
	keyPrefix    string
	localLimiter *LocalRateLimiter // Fallback for when Redis is unavailable
}

// LocalRateLimiter provides in-memory rate limiting as fallback
type LocalRateLimiter struct {
	clients map[string]*ClientLimiter
	mutex   sync.RWMutex
	cleanup time.Duration
}

// ClientLimiter tracks individual client rate limits
type ClientLimiter struct {
	tokens     int64
	lastRefill time.Time
	mutex      sync.Mutex
}

// NewRateLimiter creates a new rate limiter with Redis backend
func NewRateLimiter(windowSize time.Duration, maxRequests int64) *RateLimiter {
	return &RateLimiter{
		redis:       providers.Redis,
		windowSize:  windowSize,
		maxRequests: maxRequests,
		keyPrefix:   "rate_limit:",
		localLimiter: &LocalRateLimiter{
			clients: make(map[string]*ClientLimiter),
			cleanup: 5 * time.Minute,
		},
	}
}

// LoginRateLimit provides specialized rate limiting for login attempts
func LoginRateLimit() gin.HandlerFunc {
	limiter := NewRateLimiter(15*time.Minute, 5) // 5 attempts per 15 minutes
	
	return func(c *gin.Context) {
		clientIP := c.ClientIP()
		key := fmt.Sprintf("login_attempts:%s", clientIP)
		
		allowed, remaining, resetTime, err := limiter.Allow(key, clientIP)
		if err != nil {
			// Log error but don't block request on rate limiter failure
			c.Next()
			return
		}
		
		if !allowed {
			c.Header("X-RateLimit-Limit", strconv.FormatInt(limiter.maxRequests, 10))
			c.Header("X-RateLimit-Remaining", "0")
			c.Header("X-RateLimit-Reset", strconv.FormatInt(resetTime.Unix(), 10))
			c.Header("Retry-After", strconv.FormatInt(int64(limiter.windowSize.Seconds()), 10))
			
			c.JSON(http.StatusTooManyRequests, bases.ErrorResponse{
				ErrorCode: http.StatusTooManyRequests,
				ErrorMsg:  "Too many login attempts. Please try again later.",
				RequestID: c.GetString("RequestID"),
			})
			c.Abort()
			return
		}
		
		// Set rate limit headers
		c.Header("X-RateLimit-Limit", strconv.FormatInt(limiter.maxRequests, 10))
		c.Header("X-RateLimit-Remaining", strconv.FormatInt(remaining, 10))
		c.Header("X-RateLimit-Reset", strconv.FormatInt(resetTime.Unix(), 10))
		
		c.Next()
	}
}

// APIRateLimit provides general API rate limiting
func APIRateLimit() gin.HandlerFunc {
	limiter := NewRateLimiter(1*time.Minute, 60) // 60 requests per minute
	
	return func(c *gin.Context) {
		clientIP := c.ClientIP()
		userAgent := c.GetHeader("User-Agent")
		
		// Create composite key for more specific rate limiting
		key := fmt.Sprintf("api_requests:%s", clientIP)
		
		allowed, remaining, resetTime, err := limiter.Allow(key, clientIP)
		if err != nil {
			// Continue on rate limiter errors, but log them
			c.Next()
			return
		}
		
		if !allowed {
			c.Header("X-RateLimit-Limit", strconv.FormatInt(limiter.maxRequests, 10))
			c.Header("X-RateLimit-Remaining", "0")
			c.Header("X-RateLimit-Reset", strconv.FormatInt(resetTime.Unix(), 10))
			c.Header("Retry-After", "60")
			
			// Log rate limit violations for monitoring
			logRateLimitViolation(clientIP, userAgent, c.Request.URL.Path)
			
			c.JSON(http.StatusTooManyRequests, bases.ErrorResponse{
				ErrorCode: http.StatusTooManyRequests,
				ErrorMsg:  "Rate limit exceeded. Please slow down your requests.",
				RequestID: c.GetString("RequestID"),
			})
			c.Abort()
			return
		}
		
		// Set informational headers
		c.Header("X-RateLimit-Limit", strconv.FormatInt(limiter.maxRequests, 10))
		c.Header("X-RateLimit-Remaining", strconv.FormatInt(remaining, 10))
		c.Header("X-RateLimit-Reset", strconv.FormatInt(resetTime.Unix(), 10))
		
		c.Next()
	}
}

// Allow checks if a request should be allowed based on rate limits
func (rl *RateLimiter) Allow(key, clientIP string) (allowed bool, remaining int64, resetTime time.Time, err error) {
	redisKey := rl.keyPrefix + key
	now := time.Now()
	windowStart := now.Truncate(rl.windowSize)
	resetTime = windowStart.Add(rl.windowSize)
	
	// Try Redis first
	if rl.redis != nil {
		count, err := rl.redis.Incr(redisKey).Result()
		if err == nil {
			// Set expiry on first request in window
			if count == 1 {
				rl.redis.Expire(redisKey, rl.windowSize)
			}
			
			remaining = rl.maxRequests - count
			if remaining < 0 {
				remaining = 0
			}
			
			return count <= rl.maxRequests, remaining, resetTime, nil
		}
		// Fall through to local limiter on Redis error
	}
	
	// Fallback to local rate limiter
	return rl.localLimiter.Allow(clientIP, rl.maxRequests, rl.windowSize)
}

// Allow implements token bucket algorithm for local rate limiting
func (ll *LocalRateLimiter) Allow(clientIP string, maxRequests int64, windowSize time.Duration) (bool, int64, time.Time, error) {
	ll.mutex.Lock()
	defer ll.mutex.Unlock()
	
	now := time.Now()
	resetTime := now.Truncate(windowSize).Add(windowSize)
	
	client, exists := ll.clients[clientIP]
	if !exists {
		client = &ClientLimiter{
			tokens:     maxRequests - 1, // Consume one token for this request
			lastRefill: now,
		}
		ll.clients[clientIP] = client
		return true, maxRequests - 1, resetTime, nil
	}
	
	client.mutex.Lock()
	defer client.mutex.Unlock()
	
	// Refill tokens based on time passed
	timePassed := now.Sub(client.lastRefill)
	if timePassed >= windowSize {
		client.tokens = maxRequests
		client.lastRefill = now
	}
	
	if client.tokens > 0 {
		client.tokens--
		return true, client.tokens, resetTime, nil
	}
	
	return false, 0, resetTime, nil
}

// StartCleanup starts background cleanup of expired entries
func (ll *LocalRateLimiter) StartCleanup() {
	go func() {
		ticker := time.NewTicker(ll.cleanup)
		defer ticker.Stop()
		
		for range ticker.C {
			ll.cleanup_expired()
		}
	}()
}

// cleanup_expired removes expired client limiters
func (ll *LocalRateLimiter) cleanup_expired() {
	ll.mutex.Lock()
	defer ll.mutex.Unlock()
	
	now := time.Now()
	for ip, client := range ll.clients {
		client.mutex.Lock()
		if now.Sub(client.lastRefill) > ll.cleanup {
			delete(ll.clients, ip)
		}
		client.mutex.Unlock()
	}
}

// logRateLimitViolation logs rate limit violations for security monitoring
func logRateLimitViolation(clientIP, userAgent, path string) {
	// This should integrate with your logging system
	// Consider sending alerts for repeated violations
	fmt.Printf("RATE_LIMIT_VIOLATION: IP=%s UserAgent=%s Path=%s Time=%s\n", 
		clientIP, userAgent, path, time.Now().Format(time.RFC3339))
}